#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
模型下载工具 (使用镜像站点)

使用方法:
    python download_model.py     # 下载项目所需的两个模型
"""

import os
from huggingface_hub import snapshot_download

# 设置镜像站点
os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com")

# 项目所需的全部模型
REQUIRED_MODELS = [
    "Salesforce/codet5-base",   # 基础训练、多任务学习
    "Salesforce/codet5-small",  # 知识蒸馏学生模型 / 轻量化推理
]

def download_model(model_name: str = "Salesforce/codet5-small"):
    """下载指定的 Hugging Face 模型到本地"""
    local_dir = f"./models/{model_name.replace('/', '_')}"
    
    # 检查是否已下载
    if os.path.exists(local_dir) and os.path.isdir(local_dir):
        print(f"[跳过] 模型已存在: {local_dir}")
        return local_dir
    
    print(f"\n{'='*60}")
    print(f"[INFO] 开始下载模型: {model_name}")
    print(f"[INFO] 使用镜像站点: https://hf-mirror.com")
    print(f"[INFO] 本地保存路径: {local_dir}")
    print(f"{'='*60}\n")
    
    try:
        snapshot_download(
            repo_id=model_name,
            local_dir=local_dir,
            local_dir_use_symlinks=False,
            endpoint=os.environ.get("HF_ENDPOINT", "https://hf-mirror.com")
        )
        print(f"\n[成功] 模型下载完成: {model_name}\n")
        return local_dir
    except Exception as e:
        print(f"\n[错误] 下载失败: {e}\n")
        return None

def download_all_models():
    """下载项目所需的全部模型"""
    print("\n" + "="*60)
    print("Text2Code 项目模型下载")
    print("="*60)
    print(f"\n需要下载 {len(REQUIRED_MODELS)} 个模型:")
    for i, model in enumerate(REQUIRED_MODELS, 1):
        print(f"  {i}. {model}")
    print("")
    
    success_count = 0
    failed_models = []
    
    for model in REQUIRED_MODELS:
        result = download_model(model)
        if result:
            success_count += 1
        else:
            failed_models.append(model)
    
    print("\n" + "="*60)
    print("下载汇总")
    print("="*60)
    print(f"成功: {success_count}/{len(REQUIRED_MODELS)}")
    
    if failed_models:
        print(f"失败: {len(failed_models)}")
        for model in failed_models:
            print(f"  - {model}")
    else:
        print("\n[成功] 所有模型已下载完成!")
    print("")

if __name__ == "__main__":
    # 固定行为: 始终下载项目所需的两个模型
    download_all_models()